% Created 2021, Mike Siedlik
% Free for research and non-commercial use
% Please cite Siedlik and Issadore, 2021

% This MATLAB code was created with the following version of MATLAB:
% MATLAB Version: 9.10.0.1739362 (R2021a) Update 5
% This code uses the following MATLAB toolboxes:
% Image Processing Toolbox

% ----------------------------------------------------------------
% STEP 1: LOAD IMAGES
cntl_data = struct;

% Specify image names for each dilution
cntl_data(1).pic_name = 'B_1_1x.tif';
cntl_data(2).pic_name = 'B_2_0.5x.tif';
cntl_data(3).pic_name = 'B_3_0.3x.tif';
cntl_data(4).pic_name = 'B_4_0.2x.tif';
cntl_data(5).pic_name = 'B_5_0.05x.tif';

% Specify the known droplet dilution. In this case, this value corresponds
% to a fraction of the undiluted dye concentration
cntl_data(1).dilution = 1;
cntl_data(2).dilution = 0.5;
cntl_data(3).dilution = 0.3;
cntl_data(4).dilution = 0.2;
cntl_data(5).dilution = 0.05;

% Expected distance between droplets:
dx = 300;

% Expected droplet length
dL = 150;

% Expected number of droplets. There will be more than 3 in each image,
% though this will be sufficient to cover most of the channel region.
numD = 3;

% Box to crop final images (optional). For the example data, this will be
% useful 
cB = [236,513];


% ----------------------------------------------------------------
% ----------------------------------------------------------------
% STEP 2: CREATE A SINGLE REFERENCE IMAGE FOR EACH DILUTION
% Loop through all dilutions
for m = 1:5

    % Obtain the useful dimensions of each image sequence: number rows,
    % number columns, number frames
    INFO = imfinfo(cntl_data(m).pic_name);
    im_dim = [INFO(1).Height,INFO(1).Width,numel(INFO)];
    
    % Preallocate 3D matrix to store grey values that correspond to
    % internal droplet regions.
    internalDropPixels = zeros(im_dim,'uint8');

    % Read reference images corresponding the the back and front of a
    % droplet. Note: in these images, the droplets are moving from left to
    % right.
    % imB: image of "back" region, type single
    % imB_d: dimensions of "back" reference image
    [imB,imB_d] = loadReferenceImage(strcat('B_',cntl_data(m).pic_name(3),...
        '_back.tif'));
    [imF,imF_d] = loadReferenceImage(strcat('B_',cntl_data(m).pic_name(3),...
        '_front.tif'));

    % Loop through all frames in the m-th dilution image sequence
    for k = 1:im_dim(3)

        % Pre-process the k-th frame
        [im_cropped,im_k] = ...
            preprocessFrame(cntl_data(m).pic_name,k,imB_d,im_dim);

        % Calculate sum square error for each cropped image. These two
        % vectors contain 
        errB = squeeze(sum((imB - im_cropped).^2,[1,2])./numel(imB));
        errF = squeeze(sum((imF - im_cropped).^2,[1,2])./numel(imF));

        % Determine indices of the first numD droplets
        d_ind = identifyDropletBounds(numD,errB,errF,dx,dL);

        % Create matrix to store indices corresponding to "internal droplet
        % regions."
        bw_Drops = zeros(im_dim(1:2),'uint8');
        
        % Consider each droplet
        for mm = 1:numD
            % Index "directly to the right" of where the image best
            % correlates with the "back" reference image.
            dStart = d_ind(mm,1) + imB_d(2);
            % Index "directly to the left" of where the image best
            % correlates with the "front" reference image.
            dEnd = d_ind(mm,2);
            % Set all pixels between the back and front of the droplet to
            % 1.
            bw_Drops(:,dStart:dEnd) = 1;
        end

        % Store the internal grey values for the k-th image frame.
        internalDropPixels(:,:,k) = bw_Drops.*im_k;
        

    end

    % Convert to single for the subsequent NaN step
    internalDropPixels = single(internalDropPixels);
    
    % Replace non-droplet regions with NaN.
    internalDropPixels(internalDropPixels == 0) = nan;

    % Take the mean of all non-NaN pixels on a pixel-by-pixel basis.
    cntl_data(m).meanGreyValues = ...
        uint8(mean(internalDropPixels(:,cB(1):cB(2),:),3,'omitnan'));
    
    % (Optional) Write the image corresponding to the average grey values
    imwrite(cntl_data(m).meanGreyValues,...
        strcat('fullIm_',cntl_data(m).pic_name));

end




function [im_ref,l_dim] = loadReferenceImage(refName)
    % Local function for to load reference image, convert to single
    % channel, and return useful dimensions
    
    % Load reference image
    im_ref = imread(refName);
    
    % Take mean of each channel
    im_ref = round(mean(im_ref,3));
    
    % Convert to type single. This is the single channel reference image.
    im_ref = single(im_ref);

    % Determine dimensions of reference image
    l_dim = size(im_ref);

end


function [im_matrix,im_8] = preprocessFrame(picName,frame,imB_d,im_dim)
    % Local function to pre-process a frame of the reference image stack.
    %
    % This function returns the following two matrices:
    % im_matrix: a 3D matrix where each page corresponds to a "cropped"
    % region of this image. The first page corresponds to the region at the
    % most-left portion of the channel, then each successive page starts
    % the cropping 1 pixel further along the channel. 
    % Each page has the same dimensions as imB and imF.
    % im_8: 8-bit image of reference image.
    
    % Read image
    im = imread(picName,frame);

    % Convert to single channel by taking mean of RGB values
    im_8 = uint8(mean(im,3));

    % Index matrix used for cropping
    ind_mat = ones(im_dim(2)-imB_d(2),1)*(1:prod(imB_d)) + ...
         ((0:im_dim(2)-imB_d(2)-1)*imB_d(1))'*ones(1,prod(imB_d));
    
    % Perform indexing -- this effectively performs many cropping
    % operations in a vectorized manner
    im2 = im_8(ind_mat');

    % Reshape, where each page is one cropped image
    im_matrix = single(reshape(im2,imB_d(1),imB_d(2),im_dim(2)-imB_d(2)));
end


function dropInd = identifyDropletBounds(numD,errB,errF,dx,dL)
    % Local function to determine the left and rightmost bounds of each
    % internal droplet region.

    % Preallocate vectors to store positions (pixel indices) of the front
    % and back of the droplets.
    I_F = zeros(1,numD);
    I_B = zeros(1,numD);

    % Determine position of the back of the first droplet.
    [~,I_B(1)] = min(errB(1:dx));
    
    % Create region to search for the front of the previously idenfied
    % droplet.
    xDDist = I_B(1):I_B(1) + dL;
    
    % Determine position of front of droplet, ensuring that the pixel
    % indices correspond to the original image frame.
    [~,i_f] = min(errF(xDDist));
    I_F(1) = i_f + I_B(1) - 1;

    % Repeat the above operations for each of the expected droplets
    for mm = 2:numD

        % Determine region to search for back of the k-th droplet. The
        % minimum function is to ensure that this doesn't search beyond the
        % number of columns in the image.
        xDLeng = I_F(mm-1):min(I_F(mm-1) + dx,numel(errB));

        % Determine position of back of the next droplet.
        [~,i_b] = min(errB(xDLeng)); 
        I_B(mm) = i_b + I_F(mm-1) - 1;     

        % Determine region to search for front of the k-th droplet. The
        % minimum function is to ensure that this doesn't search beyond the
        % number of columns in the image.
        xDDist = I_B(mm):(min(I_B(mm) + dL,numel(errF)));
        
        % Determine position of front of the next droplet.
        [~,i_f] = min(errF(xDDist));
        I_F(mm) = i_f + I_B(mm) - 1;
    end
    
    % Return matrix where the each row corresponds to a droplet, the first
    % column corresponds to the "back" of that droplet, and the second
    % column corresponds to the "front" of that droplet.
    dropInd = [I_B',I_F'];
    
end

